// Copyright (c) 2020 cgsdfc
//
// This software is released under the MIT License.
// https://opensource.org/licenses/MIT

// This file implements the deepcopy function.
#pragma once

namespace pyast {

namespace detail {

template <typename T>
T* CopyBytes(Arena* arena, T* src) {
  PYAST_INTERNAL_DCHECK(arena, "");
  PYAST_INTERNAL_DCHECK(src, "");
  void* dst = std::memcpy(arena.Malloc(sizeof(T)), src, sizeof(T));
  return reinterpret_cast<T*>(dst);
}

// Handle the actual deepcopy impl.
template <typename T>
struct DeepCopyTraits;

class DeepCopyContext {
 public:
  static constexpr int kMaxRecursionLevel = 4096;

  template <typename T>
  static T DeepCopy(const T& v, Arena* arena) {
    DeepCopyContext ctx(v, arena);
    return T(ctx.arena_, ctx.DoCopy<T>(v.ptr()));
  }

 private:
  DeepCopyContext(const AstNode& root, Arena arena)
      : arena_(arena.is_null() ? root.arena() : arena),
        cross_arena_(!arena.is_null()) {
    PYAST_INTERNAL_DCHECK(arena_, "");
  }

  // Choose between plain T or AsdlSeq<T>. See below.
  template <typename T>
  struct Helper;

  template <typename T>
  typename T::underlying_type DoCopy(void* v) {
    static_assert(IsAST<T>{}, "Non-AST must not call this");
    return reinterpret_cast<typename T::underlying_type>(
        DoCopyImpl(v, &Helper<T>::RecursivelyCopy));
  }

  using Callback = void* (*)(DeepCopyContext*, void*);

  void* DoCopyImpl(void* v, Callback callback);

  void MaybeAddObject(pybind11::handle h) {
    if (cross_arena_)
      arena->AddObject(pybind11::reinterpret_borrow<object>(h));
  }

  class Scope {
   public:
    explicit Scope(DeepCopyContext* ctx) : ctx_(ctx) {
      DCHECK_TOGGLE_EXC(ctx->level_ < kMaxRecursionLevel,
                        "Maximum recursion level reached during DeepCopy()",
                        std::runtime_error);
      ++ctx_->level_;
    }

    ~Scope() { --ctx_->level_; }

   private:
    DeepCopyContext* ctx_;
  };

  template <typename T>
  friend class DeepCopyTraits;

  Arena* arena_;
  bool cross_arena_;
  std::map<void*, void*> memo_;
  int level_ = 0;
};


template <typename T>
struct DeepCopyTraits<AsdlSeq<T>> {
  using ListType = AsdlSeq<T>;
  using Traits = typename ListType::Traits;
  using underlying_type = typename ListType::underlying_type;

  // Non-AST overload
  static void* RecursivelyCopyImpl(DeepCopyContext* ctx,
                                   underlying_type v,
                                   std::false_type) {
    return Traits::CopySeq(ctx->arena_, v);
  }

  // AST overload.
  static void* RecursivelyCopyImpl(DeepCopyContext* ctx,
                                   underlying_type v,
                                   std::true_type) {
    auto size = Traits::GetSize(v);
    underlying_type u = Traits::NewSeq(ctx->arena_, size);
    for (int i = 0; i < size; i++) {
      void* old_elm = Traits::GetItem(v, i);
      void* new_elm = ctx->DoCopy<T>(old_elm);
      Traits::SetItem(u, i, new_elm);
    }
    return u;
  }

  static void* RecursivelyCopy(DeepCopyContext* ctx, void* v) {
    return RecursivelyCopyImpl(ctx, reinterpret_cast<underlying_type>(v),
                               IsAST<T>{});
  }
};

// Default is for non-AsdlSeq type that uses the underlying_type.
// These Traits specializations are generated.
template <typename T>
struct DeepCopyContext::Helper : DeepCopyTraits<typename T::underlying_type> {};

// For AsdlSeq-type, a specialized handler.
template <typename T>
struct DeepCopyContext::Helper<AsdlSeq<T>> : DeepCopyTraits<AsdlSeq<T>> {};

}  // namespace detail

template <typename T, PYAST_ENABLE_IF(detail::IsAST<T>{})>
T Copy(const T& v, Arena* arena = nullptr) {
  if (!v)
    return v;
  arena = arena ? arena : v.GetArena();
  return T(arena, detail::CopyBytes(arena, v.ptr()));
}

template <typename T>
AsdlSeq<T> Copy(const AsdlSeq<T>& v, Arena arena = nullptr) {
  return AsdlSeq<T>::Copy(v, arena);
}

template <typename T>
T DeepCopy(const T& v, Arena arena = nullptr) {
  return v.is_null() ? v : detail::DeepCopyContext::DeepCopy(v, arena);
}

}  // namespace pyast
